#include <chrono>
#include <exception>
#include <iomanip>

#include "config.hpp"
#include "header.hpp"
#include "logger.hpp"
#include "raw_save.hpp"
#include "utils.hpp"

/**
 * @brief 构造函数
 */
raw_save::raw_save()
{
}

/**
 * @brief 构造函数
 * @param ringbuffer 环形缓冲区，缓存 H264 文件
 * @param ip 相机 IP 地址
 */
raw_save::raw_save(ringbuffer* ringbuffer, std::string ip)
{
    if (nullptr == ringbuffer) {
        return;
    }

    this->ip = ip;
    this->rbuf = ringbuffer;
    utils::folder_create(RAW_VIDEO_SAVE_PATH + ip, true);
}

/**
 * @brief 析构函数
 */
raw_save::~raw_save()
{
}

/**
 * @brief 设置对象参数
 * @param ringbuffer 环形缓冲区，缓存 H264 文件
 * @param ip 相机 IP 地址
 */
int raw_save::set(ringbuffer* ringbuffer, std::string ip)
{
    if (nullptr == ringbuffer) {
        return -1;
    }

    this->ip = ip;
    this->rbuf = ringbuffer;
    utils::folder_create(RAW_VIDEO_SAVE_PATH + ip, true);

    return 0;
}

/**
 * @brief 主线程，用于保存原始视频，并封装成 MP4 文件
 */
void raw_save::thread()
{
    int ret = 0;
    int cnt = 0;
    int ridx = 0;
    int64_t first_nalu_pts = 0;
    int64_t first_nalu_dts = 0;
    bool recording = false;
    header_t head = { 0 };
    AVPacket packet = {};
    std::chrono::system_clock::time_point start_time;

    if (nullptr == this->rbuf) {
        return;
    }

    int max_size = 2 * 1024 * 1024;
    uint8_t* nalu = (uint8_t*)av_malloc(max_size);

    while (running) {
        /* 逐帧读取环形缓冲区视频帧 */
        rbuf->lock();

        /* 先读取信息头 */
        int curr = ridx;
        ret = rbuf->readr(&ridx, (uint8_t*)&head, sizeof(header_t));
        if (ret) {
            rbuf->unlock();
            msleep(40);
            continue;
        }

        /* 判断信息头同步码 */
        if (head.sync[0] != 'H' || head.sync[1] != 'E' || head.sync[2] != 'A' || head.sync[3] != 'D') {
            rbuf->unlock();
            Logger::error("Camera {} raw, Sync code not match", ip);
            msleep(40);
            std::terminate();
            continue;
        }

        /* 手动计算时间戳 */
        memset(nalu, 0, max_size);
        packet.data = nalu;
        packet.size = head.size;
        packet.pts = head.pts;
        packet.dts = head.dts;
        packet.flags = head.type;
        ret = rbuf->readr(&ridx, packet.data, packet.size);
        if (ret) {
            rbuf->unlock();
            msleep(40);
            continue;
        }
        rbuf->unlock();

        /* 只有成功拿到 H.264 码流才会打印，并且打印的是当前帧的缓存区位置 */
        // Logger::debug("Camera {} [{}] raw, read index = {}", ip, cnt, curr);
        cnt++;

        if (!recording && H264_IDR == head.type) {
            /* 刚开机，需要等 IDR 帧才会开始存原始视频 */

            /* 等到了第一个 IDR 帧，创建第一个 MP4 文件 */
            std::string name = RAW_VIDEO_SAVE_PATH + ip + "/" + utils::time_get() + ".mp4";
            ret = open(name);
            if (ret) {
                Logger::error("Camera {} raw, Failed to open mp4 file", ip);
                recording = false;
                msleep(40);
                continue;
            }

            /* 记录当前 MP4 第一帧视频的时间戳 */
            first_nalu_pts = head.pts;
            first_nalu_dts = head.dts;

            /* 文件创建好后，将 IDR 帧写入 MP4 文件中 */
            /* MP4 文件的时间戳不使用 rtsp 时间戳，而是从 0 开始 */
            packet.pts = head.pts - first_nalu_pts;
            packet.dts = head.dts - first_nalu_dts;
            ret = write(&packet);
            if (ret) {
                Logger::error("Camera {} raw, Failed to write mp4 file", ip);
                recording = false;
                msleep(40);
                continue;
            }

            /* 正在录制标记置位 */
            recording = true;

            /* 记录开始录制的起始时间 */
            start_time = std::chrono::system_clock::now();
        } else {
            /**
             * 正在录制有两种情况：
             * 第一种情况是不足 RAW_VIDEO_TIME 秒，则持续写入视频帧到 MP4 文件中，直到 MP4 文件时长满足要求
             * 第二种情况是时间满足要求了，这个时候，把所有的 P 帧都写完，等到下一个 IDR 再重新创建一个新的文件
             */

            auto now = std::chrono::system_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::seconds>(now - start_time).count();

            if (duration < RAW_VIDEO_TIME) {
                /* 视频时长不够，持续录制，直到时长满足要求 */
                packet.pts = head.pts - first_nalu_pts;
                packet.dts = head.dts - first_nalu_dts;
                ret = write(&packet);
                if (ret) {
                    Logger::error("Camera {} raw, Failed to write mp4 file", ip);
                    recording = false;
                    msleep(40);
                    continue;
                }
            } else {
                if (H264_IDR != head.type) {
                    /* 时长满足了，但是由于不是 IDR 帧，因此不能在这个位置切片，不然会导致后面的 P 帧丢失 */
                    packet.pts = head.pts - first_nalu_pts;
                    packet.dts = head.dts - first_nalu_dts;
                    ret = write(&packet);
                    if (ret) {
                        Logger::error("Camera {} raw, Failed to write mp4 file", ip);
                        recording = false;
                        msleep(40);
                        continue;
                    }
                } else {
                    /* 时长满足了，并且此时是 IDR 帧，则在这个位置切片 */

                    /* 关闭 MP4 文件 */
                    close();

                    /* 创建新的 MP4 文件 */
                    std::string name = RAW_VIDEO_SAVE_PATH + ip + "/" + utils::time_get() + ".mp4";
                    ret = open(name);
                    if (ret) {
                        Logger::error("Camera {} raw, Failed to open mp4 file", ip);
                        recording = false;
                        msleep(40);
                        continue;
                    }

                    /* 记录当前 MP4 第一帧视频的时间戳 */
                    first_nalu_pts = head.pts;
                    first_nalu_dts = head.dts;

                    /* pts 从 0 开始 */
                    packet.pts = head.pts - first_nalu_pts;
                    packet.dts = head.dts - first_nalu_dts;

                    /* 文件创建好后，将 IDR 帧写入 MP4 文件中 */
                    ret = write(&packet);
                    if (ret) {
                        Logger::error("Camera {} raw, Failed to write mp4 file", ip);
                        recording = false;
                        msleep(40);
                        continue;
                    }

                    /* 正在录制标记置位 */
                    recording = true;

                    /* 记录开始录制的起始时间 */
                    start_time = std::chrono::system_clock::now();
                }
            }
        }
    }

    av_free(nalu);

    /* 被 stop() 接口退出本线程，则主动关闭 MP4 文件，避免出现 .tmp 文件 */
    if (recording) {
        close();
    }
}

/**
 * @brief 启动原始视频保存线程
 */
void raw_save::start()
{
    if (!handle.joinable()) {
        running = true;
        handle = std::thread(&raw_save::thread, this);
    }
}

/**
 * @brief 关闭原始视频保存线程
 */
void raw_save::stop()
{
    if (handle.joinable()) {
        running = false;
        handle.join();
    }
}

/**
 * @brief 创建 MP4 文件
 */
int raw_save::open(std::string filename)
{
    this->filename = filename;
    std::string tempname = filename + ".tmp";
    Logger::debug("Camera {} raw, file create {}", ip, filename);

    /* 明确指定输出格式为 MP4 */
    AVOutputFormat* ofmt = av_guess_format("mp4", nullptr, nullptr);
    if (!ofmt) {
        Logger::error("Camera {} raw, Could not find output format for MP4", ip);
        return -1;
    }

    if (avformat_alloc_output_context2(&fmtctx, ofmt, nullptr, tempname.c_str()) < 0) {
        Logger::error("Camera {} raw, Could not create output context", ip);
        return -1;
    }

    codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        Logger::error("Camera {} raw, H.264 codec not found", ip);
        return -1;
    }

    stream = avformat_new_stream(fmtctx, codec);
    if (!stream) {
        Logger::error("Camera {} raw, Failed to create video stream", ip);
        return -1;
    }

    AVCodecParameters* codecpar = stream->codecpar;
    codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    codecpar->codec_id = AV_CODEC_ID_H264;
    codecpar->width = 1920;
    codecpar->height = 1080;
    codecpar->format = AV_PIX_FMT_YUV420P;

    if (!(fmtctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&(fmtctx->pb), tempname.c_str(), AVIO_FLAG_WRITE) < 0) {
            Logger::error("Camera {} raw, Could not open output file", ip);
            return -1;
        }
    }

    if (avformat_write_header(fmtctx, nullptr) < 0) {
        Logger::error("Camera {} raw, Error writing header to output file", ip);
        return -1;
    }

    return 0;
}

/**
 * @brief 写 MP4 文件
 */
int raw_save::write(AVPacket* packet)
{
    if (av_interleaved_write_frame(fmtctx, packet) < 0) {
        Logger::error("Camera {} raw, Error writing frame", ip);
        return -1;
    }

    return 0;
}

/**
 * @brief 关闭 MP4 文件
 */
int raw_save::close()
{
    av_write_trailer(fmtctx);
    avio_closep(&(fmtctx->pb));

    // tmp 文件重命名为 mp4 文件
    const char* finalname = filename.c_str();
    if (rename((filename + ".tmp").c_str(), finalname) != 0) {
        Logger::error("Camera {} raw, Error renaming file to {}", ip, finalname);
        return -1;
    }
    Logger::debug("Camera {} raw, file rename {}", ip, finalname);

    avformat_free_context(fmtctx);

    return 0;
}

/**
 * @brief 毫秒延时
 */
void raw_save::msleep(int ms)
{
    std::this_thread::sleep_for(std::chrono::milliseconds(ms));
}
